#!/usr/bin/env julia

"""
Convergence Study Example

This example performs a comprehensive mesh refinement study to demonstrate
convergence rates and grid independence for the 2D lid-driven cavity problem.

Run with: julia convergence_study.jl

Features:
- Mesh refinement study with automatic grid generation
- Convergence rate analysis using Richardson extrapolation
- Grid independence testing
- Spatial discretization error quantification
- Adaptive mesh refinement recommendations
- Comparison of different discretization orders
- Solution accuracy assessment at various grid resolutions

"""

# Add the package to the path
push!(LOAD_PATH, "../src")

using LidDrivenCavity
using Printf
using Statistics
using LinearAlgebra

"""
    compute_solution_norms(result::LidDrivenResult) -> NamedTuple

Compute various solution norms for convergence analysis.
"""
function compute_solution_norms(result::LidDrivenResult)
    # Extract fields and compute norms on fluid region
    ψ = result.ψ
    ω = result.ω
    u = result.u
    v = result.v
    
    # Filter out NaN values (non-fluid regions)
    ψ_fluid = ψ[.!isnan.(ψ)]
    ω_fluid = ω[.!isnan.(ω)]
    u_fluid = u[.!isnan.(u)]
    v_fluid = v[.!isnan.(v)]
    
    return (
        psi_l2 = norm(ψ_fluid, 2),
        psi_inf = norm(ψ_fluid, Inf),
        omega_l2 = norm(ω_fluid, 2),
        omega_inf = norm(ω_fluid, Inf),
        u_l2 = norm(u_fluid, 2),
        u_inf = norm(u_fluid, Inf),
        v_l2 = norm(v_fluid, 2),
        v_inf = norm(v_fluid, Inf),
        velocity_l2 = sqrt(norm(u_fluid, 2)^2 + norm(v_fluid, 2)^2),
        n_points = length(ψ_fluid)
    )
end

"""
    extract_centerline_profiles(result::LidDrivenResult) -> NamedTuple

Extract velocity profiles along cavity centerlines for detailed comparison.
"""
function extract_centerline_profiles(result::LidDrivenResult)
    x, y = result.x, result.y
    u, v = result.u, result.v
    n = length(x)
    
    # Find centerline indices (approximately)
    center_i = div(n, 2)
    center_j = div(n, 2)
    
    # Extract profiles avoiding NaN regions
    u_vertical = Float64[]    # u along vertical centerline
    v_horizontal = Float64[]  # v along horizontal centerline
    y_vertical = Float64[]
    x_horizontal = Float64[]
    
    # Vertical centerline profile (u-velocity)
    for j in 1:n
        if !isnan(u[center_i, j])
            push!(u_vertical, u[center_i, j])
            push!(y_vertical, y[j])
        end
    end
    
    # Horizontal centerline profile (v-velocity)
    for i in 1:n
        if !isnan(v[i, center_j])
            push!(v_horizontal, v[i, center_j])
            push!(x_horizontal, x[i])
        end
    end
    
    return (
        u_vertical = u_vertical,
        y_vertical = y_vertical,
        v_horizontal = v_horizontal,
        x_horizontal = x_horizontal
    )
end

"""
    compute_convergence_rate(h_values, error_values) -> Float64

Compute convergence rate using least squares fit of log(error) vs log(h).
"""
function compute_convergence_rate(h_values, error_values)
    if length(h_values) < 2 || any(error_values .<= 0)
        return NaN
    end
    
    log_h = log.(h_values)
    log_error = log.(error_values)
    
    # Linear regression: log(error) = p * log(h) + c
    # where p is the convergence rate
    n = length(log_h)
    sum_x = sum(log_h)
    sum_y = sum(log_error)
    sum_xy = sum(log_h .* log_error)
    sum_x2 = sum(log_h .^ 2)
    
    # Slope is convergence rate
    rate = (n * sum_xy - sum_x * sum_y) / (n * sum_x2 - sum_x^2)
    
    return rate
end

"""
    richardson_extrapolation(coarse_sol, fine_sol, refinement_ratio=2) -> Float64

Estimate discretization error using Richardson extrapolation.
"""
function richardson_extrapolation(coarse_val, fine_val, ultra_fine_val, p=2.0)
    # Richardson extrapolation formula for error estimation
    # Assumes refinement ratio of 2
    if refinement_ratio = 2.0
        # Richardson estimate of exact solution
        exact_estimate = fine_val + (fine_val - coarse_val) / (2^p - 1)
        
        # Error estimates
        coarse_error = abs(exact_estimate - coarse_val)
        fine_error = abs(exact_estimate - fine_val)
        
        return (
            exact_estimate = exact_estimate,
            coarse_error = coarse_error,
            fine_error = fine_error,
            observed_rate = log(coarse_error / fine_error) / log(2.0)
        )
    else
        return nothing
    end
end

"""
    mesh_refinement_study(reynolds::Float64; grid_sizes=[8, 16, 32, 64])

Perform comprehensive mesh refinement study.
"""
function mesh_refinement_study(reynolds::Float64; grid_sizes=[8, 16, 32, 64])
    println("\n🔬 Mesh Refinement Study")
    println("="^50)
    println("Reynolds number: $(reynolds)")
    println("Grid sizes: $(grid_sizes)")
    println()
    
    # Storage for results
    solutions = []
    solution_norms = []
    centerline_profiles = []
    h_values = Float64[]
    solve_times = Float64[]
    iterations = Int[]
    
    # Solve on each grid
    for n in grid_sizes
        println("📐 Solving on $(n)×$(n) grid:")
        
        options = LidDrivenOptions(
            n = n,
            Re = reynolds,
            solver = :julia,  # Use consistent solver for comparison
            preconditioner = :diagonal,
            max_steps = 2000,
            tol = 1e-7,  # Tight tolerance for accuracy
            verbose = false,
            save_history = true
        )
        
        print("  Setting up and solving... ")
        flush(stdout)
        
        try
            result = solve_lid_driven_2d(options)
            
            if result.converged
                push!(solutions, result)
                push!(solution_norms, compute_solution_norms(result))
                push!(centerline_profiles, extract_centerline_profiles(result))
                push!(h_values, 2.0 / n)  # Grid spacing
                push!(solve_times, result.solve_time)
                push!(iterations, result.iterations)
                
                println("✅ ($(@sprintf("%.3f", result.solve_time))s, $(result.iterations) iter)")
            else
                println("❌ Failed to converge")
                break
            end
            
        catch e
            println("❌ Error: $e")
            break
        end
    end
    
    n_solutions = length(solutions)
    
    if n_solutions < 2
        println("❌ Need at least 2 successful solutions for convergence analysis")
        return nothing
    end
    
    println("\n📊 Grid Convergence Analysis:")
    println("Grid Size   h        DOF      Time(s)   Iterations   ψ L2-norm    ω L2-norm")
    println("-"^75)
    
    for i in 1:n_solutions
        n = grid_sizes[i]
        h = h_values[i]
        dof = solution_norms[i].n_points
        time = solve_times[i]
        iter = iterations[i]
        psi_norm = solution_norms[i].psi_l2
        omega_norm = solution_norms[i].omega_l2
        
        println(@sprintf("%4d×%-4d %6.4f   %6d   %7.3f   %8d     %9.4e  %9.4e",
                n, n, h, dof, time, iter, psi_norm, omega_norm))
    end
    
    # Convergence rate analysis
    if n_solutions >= 3
        println("\n📈 Convergence Rate Analysis:")
        
        # Compute convergence rates for different quantities
        psi_l2_norms = [s.psi_l2 for s in solution_norms]
        omega_l2_norms = [s.omega_l2 for s in solution_norms]
        u_l2_norms = [s.u_l2 for s in solution_norms]
        velocity_l2_norms = [s.velocity_l2 for s in solution_norms]
        
        # Use differences between consecutive solutions as error estimates
        psi_errors = Float64[]
        omega_errors = Float64[]
        u_errors = Float64[]
        vel_errors = Float64[]
        error_h_values = Float64[]
        
        for i in 2:n_solutions
            # Estimate error as difference from finest available solution
            finest_idx = n_solutions
            push!(psi_errors, abs(psi_l2_norms[i] - psi_l2_norms[finest_idx]))
            push!(omega_errors, abs(omega_l2_norms[i] - omega_l2_norms[finest_idx]))
            push!(u_errors, abs(u_l2_norms[i] - u_l2_norms[finest_idx]))
            push!(vel_errors, abs(velocity_l2_norms[i] - velocity_l2_norms[finest_idx]))
            push!(error_h_values, h_values[i])
        end
        
        # Filter out zero errors
        valid_indices = findall(x -> x > 1e-15, psi_errors)
        
        if length(valid_indices) >= 2
            psi_rate = compute_convergence_rate(error_h_values[valid_indices], psi_errors[valid_indices])
            omega_rate = compute_convergence_rate(error_h_values[valid_indices], omega_errors[valid_indices])
            u_rate = compute_convergence_rate(error_h_values[valid_indices], u_errors[valid_indices])
            vel_rate = compute_convergence_rate(error_h_values[valid_indices], vel_errors[valid_indices])
            
            println("Quantity           Convergence Rate   Expected Rate")
            println("-"^50)
            println(@sprintf("Stream function    %8.2f         %8.1f", psi_rate, 2.0))
            println(@sprintf("Vorticity          %8.2f         %8.1f", omega_rate, 2.0))
            println(@sprintf("u-velocity         %8.2f         %8.1f", u_rate, 2.0))
            println(@sprintf("Velocity magnitude %8.2f         %8.1f", vel_rate, 2.0))
            
            println("\n📋 Convergence Assessment:")
            rates = [psi_rate, omega_rate, u_rate, vel_rate]
            avg_rate = mean(filter(!isnan, rates))
            
            if avg_rate > 1.8
                println("✅ Excellent convergence (rate ≈ $(@sprintf("%.1f", avg_rate)))")
                println("   Grid convergence is well-established")
            elseif avg_rate > 1.2
                println("✅ Good convergence (rate ≈ $(@sprintf("%.1f", avg_rate)))")
                println("   Solutions are converging as expected")
            elseif avg_rate > 0.8
                println("⚠️  Moderate convergence (rate ≈ $(@sprintf("%.1f", avg_rate)))")
                println("   Consider finer grids or check numerical stability")
            else
                println("❌ Poor convergence (rate ≈ $(@sprintf("%.1f", avg_rate)))")
                println("   Check solver parameters and boundary conditions")
            end
        end
    end
    
    # Grid independence assessment
    if n_solutions >= 3
        println("\n🎯 Grid Independence Analysis:")
        
        # Compare solutions between finest grids
        fine_idx = n_solutions
        medium_idx = n_solutions - 1
        
        fine_norm = solution_norms[fine_idx]
        medium_norm = solution_norms[medium_idx]
        
        # Relative differences
        psi_rel_diff = abs(fine_norm.psi_l2 - medium_norm.psi_l2) / fine_norm.psi_l2 * 100
        omega_rel_diff = abs(fine_norm.omega_l2 - medium_norm.omega_l2) / fine_norm.omega_l2 * 100
        u_rel_diff = abs(fine_norm.u_l2 - medium_norm.u_l2) / fine_norm.u_l2 * 100
        
        println("Relative difference between finest two grids:")
        println(@sprintf("  Stream function:  %.2f%%", psi_rel_diff))
        println(@sprintf("  Vorticity:        %.2f%%", omega_rel_diff))
        println(@sprintf("  u-velocity:       %.2f%%", u_rel_diff))
        
        if max(psi_rel_diff, omega_rel_diff, u_rel_diff) < 1.0
            println("✅ Grid independence achieved (< 1% difference)")
            println("   Grid size $(grid_sizes[medium_idx])×$(grid_sizes[medium_idx]) is sufficient for this accuracy")
        elseif max(psi_rel_diff, omega_rel_diff, u_rel_diff) < 5.0
            println("⚠️  Near grid independence (< 5% difference)")
            println("   Consider using grid size $(grid_sizes[fine_idx])×$(grid_sizes[fine_idx]) for higher accuracy")
        else
            println("❌ Grid dependence still significant (> 5% difference)")
            println("   Recommend further refinement beyond $(grid_sizes[fine_idx])×$(grid_sizes[fine_idx])")
        end
    end
    
    # Performance scaling analysis
    if n_solutions >= 2
        println("\n⚡ Performance Scaling:")
        println("Grid Size   Time(s)   Time/DOF   Scaling Efficiency")
        println("-"^50)
        
        base_time_per_dof = solve_times[1] / solution_norms[1].n_points
        
        for i in 1:n_solutions
            time_per_dof = solve_times[i] / solution_norms[i].n_points
            efficiency = base_time_per_dof / time_per_dof * 100
            
            println(@sprintf("%4d×%-4d %7.3f   %8.2e   %8.1f%%",
                    grid_sizes[i], grid_sizes[i], solve_times[i], time_per_dof, efficiency))
        end
    end
    
    return (
        solutions = solutions,
        solution_norms = solution_norms,
        h_values = h_values,
        grid_sizes = grid_sizes[1:n_solutions]
    )
end

"""
    centerline_convergence_analysis(study_results)

Analyze convergence of centerline velocity profiles.
"""
function centerline_convergence_analysis(study_results)
    if study_results === nothing || length(study_results.solutions) < 2
        return
    end
    
    println("\n📏 Centerline Profile Convergence:")
    println("="^40)
    
    solutions = study_results.solutions
    grid_sizes = study_results.grid_sizes
    
    # Extract centerline data for finest grid
    finest_profile = extract_centerline_profiles(solutions[end])
    
    println("Centerline profile analysis:")
    println("  Vertical centerline u-velocity points: $(length(finest_profile.u_vertical))")
    println("  Horizontal centerline v-velocity points: $(length(finest_profile.v_horizontal))")
    
    # Compare profiles between different grid resolutions
    if length(solutions) >= 3
        coarse_profile = extract_centerline_profiles(solutions[end-2])
        medium_profile = extract_centerline_profiles(solutions[end-1])
        fine_profile = extract_centerline_profiles(solutions[end])
        
        # Compute profile differences (using interpolation for different grid sizes)
        if length(fine_profile.u_vertical) >= 3 && length(medium_profile.u_vertical) >= 3
            # Simple comparison at common points
            n_compare = min(length(fine_profile.u_vertical), length(medium_profile.u_vertical))
            u_profile_diff = 0.0
            
            for i in 1:3:n_compare  # Sample every 3rd point
                u_diff = abs(fine_profile.u_vertical[i] - medium_profile.u_vertical[i])
                u_profile_diff = max(u_profile_diff, u_diff)
            end
            
            println(@sprintf("  Max u-profile difference (fine vs medium): %.4f", u_profile_diff))
            
            if u_profile_diff < 0.01
                println("✅ Centerline profiles well-converged")
            elseif u_profile_diff < 0.05
                println("⚠️  Centerline profiles moderately converged")
            else
                println("❌ Centerline profiles not yet converged")
            end
        end
    end
end

"""
    adaptive_refinement_recommendations(study_results)

Generate recommendations for adaptive mesh refinement.
"""
function adaptive_refinement_recommendations(study_results)
    if study_results === nothing
        return
    end
    
    println("\n🎯 Adaptive Refinement Recommendations:")
    println("="^45)
    
    solutions = study_results.solutions
    grid_sizes = study_results.grid_sizes
    
    if length(solutions) >= 2
        # Analyze solution features for refinement guidance
        finest_solution = solutions[end]
        
        # Compute vorticity gradients as refinement indicator
        ω = finest_solution.ω
        n = size(ω, 1)
        
        # Simple gradient magnitude (avoiding NaN regions)
        high_gradient_regions = 0
        total_fluid_points = 0
        
        for i in 2:n-1, j in 2:n-1
            if !isnan(ω[i,j]) && !isnan(ω[i+1,j]) && !isnan(ω[i,j+1])
                total_fluid_points += 1
                
                # Approximate gradient magnitude
                dω_dx = abs(ω[i+1,j] - ω[i,j])
                dω_dy = abs(ω[i,j+1] - ω[i,j])
                grad_mag = sqrt(dω_dx^2 + dω_dy^2)
                
                # Flag high gradient regions (simple threshold)
                threshold = 2.0 * std(filter(!isnan, ω))
                if grad_mag > threshold
                    high_gradient_regions += 1
                end
            end
        end
        
        gradient_percentage = high_gradient_regions / total_fluid_points * 100
        
        println("Solution analysis on finest grid ($(grid_sizes[end])×$(grid_sizes[end])):")
        println(@sprintf("  High gradient regions: %.1f%% of domain", gradient_percentage))
        
        if gradient_percentage > 15.0
            println("📍 Recommendations:")
            println("  • High gradient concentration detected")
            println("  • Consider adaptive mesh refinement near:")
            println("    - Lid-driven boundary (high shear)")
            println("    - Corner regions (singularities)")
            println("    - Internal boundary corners")
            println("  • Target refinement factor: 2-4× in critical regions")
        elseif gradient_percentage > 5.0
            println("📍 Recommendations:")
            println("  • Moderate gradient levels detected")
            println("  • Uniform refinement may be sufficient")
            println("  • Consider 1.5-2× refinement for higher accuracy")
        else
            println("📍 Recommendations:")
            println("  • Solution appears smooth on current grid")
            println("  • Current resolution may be adequate")
            println("  • Further refinement primarily improves numerical error")
        end
        
        # Grid efficiency analysis
        finest_time = study_results.solutions[end].solve_time
        finest_dof = study_results.solutions[end].options.n^2
        
        if length(solutions) >= 2
            coarsest_time = study_results.solutions[1].solve_time
            coarsest_dof = study_results.solutions[1].options.n^2
            
            time_scaling = (finest_time / coarsest_time) / (finest_dof / coarsest_dof)
            
            println("\n⚖️  Computational Efficiency:")
            println(@sprintf("  Time scaling factor: %.2f (vs linear DOF scaling)", time_scaling))
            
            if time_scaling < 1.5
                println("  ✅ Excellent scaling - solver handles refinement well")
            elseif time_scaling < 3.0
                println("  ✅ Good scaling - reasonable computational cost")
            else
                println("  ⚠️  Poor scaling - consider iterative solver improvements")
            end
        end
    end
end

function main()
    println("🔬 LidDrivenCavity.jl - Convergence Study")
    println("="^50)
    
    # Display package information
    lid_driven_info()
    println()
    
    println("📋 Study Configuration:")
    println("  Focus: Mesh refinement and convergence analysis")
    println("  Method: Stream function-vorticity formulation")
    println("  Domain: L-shaped cavity")
    println()
    
    # Test parameters
    reynolds_numbers = [100.0, 400.0]
    grid_sequences = [
        [8, 16, 32, 48],      # Moderate sequence
        [16, 32, 64]          # Higher resolution (if computational resources allow)
    ]
    
    for (i, Re) in enumerate(reynolds_numbers)
        println("="^60)
        println("Study $(i): Reynolds Number = $(Re)")
        
        # Select appropriate grid sequence based on Reynolds number
        grids = i == 1 ? grid_sequences[1] : grid_sequences[2][1:2]  # Limit for higher Re
        
        try
            study_results = mesh_refinement_study(Re; grid_sizes=grids)
            
            if study_results !== nothing
                # Additional analyses
                centerline_convergence_analysis(study_results)
                adaptive_refinement_recommendations(study_results)
            else
                println("❌ Study failed - insufficient converged solutions")
            end
            
        catch e
            println("❌ Study failed with error: $e")
            if isa(e, OutOfMemoryError)
                println("💡 Try with smaller grid sizes or more memory")
            end
        end
        
        println()
    end
    
    # Final summary and guidelines
    println("="^60)
    println("📋 General Convergence Guidelines")
    println("-"^35)
    println("• Grid Independence: < 1% change between consecutive grids")
    println("• Convergence Rate: ≈ 2.0 for second-order finite differences")  
    println("• Minimum Resolution: 16×16 for qualitative results")
    println("• Recommended Resolution: 32×32 for engineering accuracy")
    println("• High-Accuracy Resolution: 64×64+ for research applications")
    println()
    println("🎯 Specific Recommendations by Reynolds Number:")
    println("  Re = 100:   32×32 grid usually sufficient")
    println("  Re = 400:   48×48 or 64×64 recommended")
    println("  Re = 1000+: 64×64+ with careful convergence monitoring")
    println()
    println("💡 Always verify grid independence for your specific application!")
end

# Run the convergence study
if abspath(PROGRAM_FILE) == @__FILE__
    try
        main()
        println("\n🎉 Convergence study completed successfully!")
        println("   Use the recommendations above to select optimal grid sizes")
    catch e
        println("\n😞 Convergence study failed: $e")
        if isa(e, LoadError) || isa(e, MethodError) 
            println("\n💡 Make sure all dependencies are available")
        elseif isa(e, OutOfMemoryError)
            println("\n💡 Reduce grid sizes or increase available memory")
        end
    end
end